import torch
import torchvision
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

transform=transforms.Compose([transforms.ToTensor()])

#训练集
trainset=torchvision.datasets.MNIST(root='./dataMNIST',
                                    train=True,
                                    download=True,
                                    transform=transform)
trainloader=torch.utils.data.DataLoader(trainset,batch_size=4,
                                        shuffle=True,num_workers=2)

#测试集
testset=torchvision.datasets.MNIST(root='./dataMNIST',
                                   train=False,
                                   download=True,
                                   transform=transform)
testloader=torch.utils.data.DataLoader(testset,batch_size=4,
                                       shuffle=False,num_workers=2)

#随机展示训练样本以确定样本导入正确
def imshow(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))

if __name__ == '__main__':
    # 选择一个 batch 的图片
    dataiter = iter(trainloader)
    images, labels = dataiter.next()

    # 显示图片
    imshow(torchvision.utils.make_grid(images))
    plt.show()
    # 打印 labels
    print(' '.join('%11s' % labels[j].numpy() for j in range(4)))

#CNN构建核心代码
class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1=nn.Conv2d(1,6,5)    #1个输入图片的通道，6个输出通道 5*5 卷积核
        self.pool=nn.MaxPool2d(2,2)    #max pool ,2×2
        self.conv2=nn.Conv2d(6,16,5)   #6个输入图片的通道，16个输出通道，5×5卷积核
        self.fc1=nn.Linear(16*4*4,120) #拉成一维向量，全连接层
        self.fc2=nn.Linear(120,84)     #全连接层
        self.fc3=nn.Linear(84,10)      #全连接层，输出层softmax ，10个数字

    def forward(self,x):
        x=self.pool(F.relu(self.conv1(x)))
        x=self.pool(F.relu(self.conv2(x)))
        x=x.view(-1,16*4*4)  #拉成一维向量
        x=F.relu(self.fc1(x))
        x=F.relu(self.fc2(x))
        x=self.fc3(x)
        return  x

net=Net()
#打印网络以确定网络结构
print(net)

# 定义损失函数，分类问题使用交叉熵，回归问题使用均方差
criterion=nn.CrossEntropyLoss()

# 构建优化器
optimizer=optim.Adam(net.parameters(),lr=0.0001)

# 训练网络
num_epoches=5 #设置 epoch数目(迭代次数
cost=[]  # 损失函数累加

if __name__ == '__main__':
    for epoch in range(num_epoches):
        running_loss=0.0
        for i,data in enumerate(trainloader,0):
            # 输入样本和标签
            inputs,labels=data

            #每次训练梯度清零
            optimizer.zero_grad()

            # 正向传播，反向传播和优化过程（网络参数更新）
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()  # 框架完成所有的反向传播和梯度计算
            optimizer.step()

            # 打印训练情况
            running_loss+=loss.item()
            if i%2000==1999: #每隔2000 mini-batches ，打印一次
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                cost.append(running_loss / 2000)
                running_loss = 0.0

    # 输出损失
    plt.plot(cost)
    plt.ylabel('cost')
    plt.show()

    # 测试数据
    correct = 0
    total = 0
    with torch.no_grad():
        for data in testloader:
            images, labels = data
            outputs = net(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    print('Accuracy of the network on the 10000 test images: %.3f %%' %
          (100 * correct / total))

